import os
from classy_vision.generic.registry_utils import import_all_packages_from_directory
from classy_vision.generic.util import load_json, load_checkpoint, update_classy_model
from classy_vision.models import build_model
import numpy as np
import cv2
import torch
import torchvision.transforms as transforms
import torch.nn.functional as F


if __name__ == '__main__':
    import_all_packages_from_directory(os.getcwd())

    config_file = '/home/lixuan/shelf_grid_detection/image_orientation_recognition/configs/scene.json'
    config = load_json(config_file)

    checkpoint_file = '/home/lixuan/shelf_grid_detection/image_orientation_recognition/output_2021-07-01T09:27:41.132533/checkpoints/model_phase-998_end.torch'
    checkpoint = load_checkpoint(checkpoint_file)

    model = build_model(config['model'])
    state_load_success = update_classy_model(
        model=model,
        model_state_dict=checkpoint['classy_state_dict']['base_model'],
        reset_heads=True,
        strict=True
    )
    model.eval()

    ori_forward = model.__class__.forward

    def warpper_forward(self, x):
        x = F.softmax(ori_forward(self, x), dim=1)
        return [x]

    model.__class__.forward = warpper_forward

    input_names = ['input%d' % i for i in range(1)]
    output_names = ['output%d' % i for i in range(1)]

    inp = np.random.rand(1, 3, 448, 448).astype(np.float32)
    inp_t = torch.FloatTensor(inp)

    with torch.no_grad():
        outputs = model(inp_t)
        outputs = [x.numpy() for x in outputs]
    print([x.shape for x in outputs])
    torch.onnx.export(
        model, (inp_t,), 'model.onnx', opset_version=9,
        input_names=input_names, output_names=output_names,
        dynamic_axes=dict([(k, {0: 'n', 2: 'h', 3: 'w'}) for k in input_names] +
                          [(k, {0: 'n'}) for k in output_names]),
        keep_initializers_as_inputs=True,
    )
